package model

import (
	"fmt"
	"github.com/spf13/viper"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"
	"strings"
	"tg-bot/common"
	"time"
)

const GET_SQL_URL = "datasource.dynamic"
const POSTGRE_SQL = "postgres://"

var USING_POSTGRE = false
var IS_MASTER_NODE = true
var USING_SQL_LITE = false
var DB = make(map[string]*gorm.DB)
var MASTER string

func SetupDB() {
	MASTER = viper.GetString("datasource.primary")
	InitDB()
}

func InitDB() {
	connectDB()
	if len(DB) == 0 {
		common.FatalLog("Connect Database number is null")
		panic("Connect Database number is null")
	}
}

func connectDB() {
	debug := viper.GetBool("debug")
	if viper.IsSet(GET_SQL_URL) {
		dsn_map := viper.GetStringMapString(GET_SQL_URL)
		for key, dsn := range dsn_map {
			var db *gorm.DB
			var err error
			if strings.HasPrefix(dsn, POSTGRE_SQL) {
				common.SysLog("using PostgreSQL as database")
				USING_POSTGRE = true
				db, err = gorm.Open(postgres.New(postgres.Config{
					DSN:                  dsn,
					PreferSimpleProtocol: true,
				}), &gorm.Config{
					PrepareStmt: true,
				})
			} else {
				db, err = gorm.Open(mysql.Open(dsn), &gorm.Config{
					PrepareStmt: true, // precompile SQL
				})
			}
			if err != nil {
				common.SysLog("Connect Database Failed, Error ==>>> " + err.Error())
				continue
			}
			sqlDB, err := db.DB()
			if err != nil {
				common.SysLog("Connect " + key + "Database Fail, Error ==>>> " + key)
				continue
			}
			sqlDB.SetMaxIdleConns(common.GetOrDefault("SQL_MAX_IDLE_CONNS", 100))
			sqlDB.SetMaxOpenConns(common.GetOrDefault("SQL_MAX_OPEN_CONNS", 1000))
			sqlDB.SetConnMaxLifetime(time.Second * time.Duration(common.GetOrDefault("SQL_MAX_LIFETIME", 60)))
			if debug {
				db = db.Debug()
			}
			DB[key] = db
			common.SysLog("Connect Database Success ==>>> " + key)
		}
	} else {
		common.SysLog("SQL_DSN not set, using SQLite as database")
		USING_SQL_LITE = true
		config := fmt.Sprintf("?_busy_timeout=%d", common.GetOrDefault("sqlite_busy_timeout", 3000))
		DB["master"], _ = gorm.Open(sqlite.Open(viper.GetString("sqlite_path")+config), &gorm.Config{
			PrepareStmt: true, // precompile SQL
		})
	}
}

func ChoiceUseDB(choice string) *gorm.DB {
	if choice == "" {
		choice = MASTER
	}
	value, err := DB[choice]
	if !err {
		common.SysLog("Database Connect Not Found, Get Map Key is " + choice)
	}
	return value
}

func CloseDB() error {
	var err error
	for key, db := range DB {
		sqlDB, err := db.DB()
		if err != nil {
			common.SysError("connect db fail" + key + "   " + err.Error())
			return err
		}
		err = sqlDB.Close()
	}
	return err
}
